//
// Created by z00850016 on 2025/8/12.
//

#ifndef LOCALVIT_NNEXECUTOR_H
#define LOCALVIT_NNEXECUTOR_H

#include "common.h"
#include "HiAiModelManagerService.h"
//#include "Tokenizer.h"

namespace localvit {

    class NNExecutor {
    public:
        NNExecutor();

        virtual ~NNExecutor();

        NNExecutor(const NNExecutor &) = delete;

        NNExecutor(NNExecutor &&) = delete;

        NNExecutor &operator=(const NNExecutor &) = delete;

        NNExecutor &operator=(NNExecutor &&) = delete;

        ModelStatus Init();

        ModelStatus InitModels(const std::string &modelPath, const std::string &modelName);

        ModelStatus DeInit();

        //ModelStatus Process(const std::string &modelName, VisionBuffer &inputBuffer, VisionBuffer &outputBuffer);
        ModelStatus Process(const std::string &modelName,
                               std::vector<VisionBuffer> &inputBuffers,
                               std::vector<VisionBuffer> &outputBuffers);


        std::vector<std::shared_ptr<hiai::AiTensor>> GetInputTensors();

        std::vector<std::shared_ptr<hiai::AiTensor>> GetOutputTensors();

        ModelStatus InitTokenizer(const std::string &vocabPath);
    private:
//        ModelStatus InitInputTensor(const std::vector<hiai::TensorDimension> &inputDimensions,
//                                       const VisionBuffer &inputBuffer);
        ModelStatus InitInputTensor(const std::vector<hiai::TensorDimension> &inputDimensions,
                                        const std::vector<VisionBuffer> &inputBuffers);

        ModelStatus InitOutputTensor(
                const std::vector<hiai::TensorDimension> &outputDimensions);
//        ModelStatus InitOutputTensor(const std::vector<hiai::TensorDimension> &outputDimensions,
//                                        std::vector<VisionBuffer> &targetVisionBuffers);
        void FillInOutputBuffer(std::vector<VisionBuffer> &outputBuffers);
    private:
        std::shared_ptr<hiai::AiModelMngerClient> modelManagerClient = nullptr;
        hiai::MemBuffer *inputMemBuffer = nullptr;
        std::shared_ptr<hiai::AiModelBuilder> modelBuilder = nullptr;
        std::vector<std::shared_ptr<hiai::AiTensor>> inputTensors;
        std::vector<std::shared_ptr<hiai::AiTensor>> outputTensors;
        //Tokenizer tokenizer;
    };
}

#endif //LOCALVIT_NNEXECUTOR_H
